package org.loong.crypto.extension.cms;

import java.io.IOException;
import java.security.PublicKey;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Primitive;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.ASN1Set;
import org.bouncycastle.asn1.ASN1TaggedObject;
import org.bouncycastle.asn1.BERSet;
import org.bouncycastle.asn1.BERTags;
import org.bouncycastle.asn1.DERSet;
import org.bouncycastle.asn1.cms.EnvelopedData;
import org.bouncycastle.asn1.cms.IssuerAndSerialNumber;
import org.bouncycastle.asn1.cms.KeyTransRecipientInfo;
import org.bouncycastle.asn1.cms.RecipientInfo;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.Certificate;
import org.bouncycastle.asn1.x509.SubjectKeyIdentifier;
import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.utils.KeyTools;

public final class CMSUtils {

    /**
     * Fetch certificate list.
     * 
     * @param certificate the certificate
     * @return the certificate in ASN.1 format
     * @throws IOException the {@link IOException}
     * @throws CertificateEncodingException the {@link CertificateEncodingException}
     */
    public static ASN1Set fetchCertificatesList(final X509Certificate certificate) throws CryptoException {
        try {
            List<ASN1Encodable> list = Collections.singletonList(Certificate.getInstance(ASN1Primitive.fromByteArray(certificate.getEncoded())));
            return createBERSetFromList(list);
        } catch (CertificateEncodingException | IOException e) {
            throw new CryptoException(e.getMessage(), e);
        }
    }

    /**
     * Fetch certificate list.
     * 
     * @param certificateChain the certificateChain
     * @return the certificate in ASN.1 format
     * @throws IOException the {@link IOException}
     * @throws CertificateEncodingException the {@link CertificateEncodingException}
     */
    public static ASN1Set fetchCertificatesList(final X509Certificate[] certificateChain) throws CryptoException {
        if (ArrayUtils.isNotEmpty(certificateChain)) {
            final List<ASN1Encodable> list = Arrays.stream(certificateChain).map(certificate -> {
                try {
                    return Certificate.getInstance(ASN1Primitive.fromByteArray(certificate.getEncoded()));
                } catch (CertificateEncodingException | IOException e) {
                    throw new CryptoException(e.getMessage(), e);
                }
            }).filter(Objects::nonNull).collect(Collectors.toList());
            return createBERSetFromList(list);
        }
        return null;
    }

    /**
     * Create ASN1Set from list.
     * 
     * @param derObjects the derObjects
     * @return the ASN1Set format
     */
    public static ASN1Set createBERSetFromList(final List<?> derObjects) {
        ASN1EncodableVector v = new ASN1EncodableVector();
        derObjects.forEach(derObj -> v.add((ASN1Encodable) derObj));
        return new BERSet(v);
    }

    /**
     * Create ASN1Set from list.
     * 
     * @param derObjects the derObjects
     * @return the ASN1Set format
     */
    public static ASN1Set createDERSetFromList(final List<?> derObjects) {
        ASN1EncodableVector v = new ASN1EncodableVector();
        derObjects.forEach(derObj -> v.add((ASN1Encodable) derObj));
        return new DERSet(v);
    }

    /**
     * Fetch wrapped data.
     * 
     * @param cmsData the cmsData
     * @return the Sequence of the data
     * @throws IOException the {@link IOException}
     */
    public static ASN1Sequence fetchWrappedData(final byte[] cmsData) throws IOException {
        final ASN1Sequence seq;
        try (final ASN1InputStream is = new ASN1InputStream(cmsData);) {
            seq = (ASN1Sequence) is.readObject();
        }

        return (ASN1Sequence) ((ASN1TaggedObject) seq.getObjectAt(1)).getBaseUniversal(true, BERTags.SEQUENCE);
    }

    /**
     * Fetch Encrypted Key Datas.
     * 
     * @param certificate the certificate
     * @param recipientInfos the recipientInfos
     * @return the encryptedKey
     * @throws CryptoException the {@link CryptoException}
     * @throws IOException the {@link IOException}
     * @throws CertificateEncodingException the {@link CertificateEncodingException}
     */
    public static Pair<AlgorithmIdentifier, byte[]> fetchEncryptedKeyDatas(final X509Certificate certificate, final Enumeration<?> recipientInfos)
            throws CryptoException, IOException, CertificateEncodingException {
        final IssuerAndSerialNumber issuerAndSerial = new IssuerAndSerialNumber(new JcaX509CertificateHolder(certificate).toASN1Structure());

        AlgorithmIdentifier keyEncryptionAlgorithm = null;
        byte[] encryptedKey = null;
        while (recipientInfos.hasMoreElements()) {
            ASN1Sequence seq = (ASN1Sequence) recipientInfos.nextElement();
            RecipientInfo reci = RecipientInfo.getInstance(seq);
            KeyTransRecipientInfo ktri = KeyTransRecipientInfo.getInstance(reci.getInfo());
            IssuerAndSerialNumber ias = IssuerAndSerialNumber.getInstance(ktri.getRecipientIdentifier().toASN1Primitive());
            if (ias.equals(issuerAndSerial)) {
                keyEncryptionAlgorithm = ktri.getKeyEncryptionAlgorithm();
                encryptedKey = ktri.getEncryptedKey().getOctets();
                break;
            }
        }

        if (keyEncryptionAlgorithm == null || encryptedKey == null) {
            throw new CryptoException("The specified user is not one of the recipients of the digital envelope.");
        }

        return Pair.of(keyEncryptionAlgorithm, encryptedKey);
    }

    /**
     * Fetch Encrypted Key Datas.
     * 
     * @param publicKey the publicKey
     * @param recipientInfos the recipientInfos
     * @return the encryptedKey
     * @throws CryptoException the {@link CryptoException}
     * @throws IOException the {@link IOException}
     * @throws CertificateEncodingException the {@link CertificateEncodingException}
     */
    public static Pair<AlgorithmIdentifier, byte[]> fetchEncryptedKeyDatas(final PublicKey publicKey, final Enumeration<?> recipientInfos)
            throws CryptoException, IOException, CertificateEncodingException {
        final byte[] keyId = KeyTools.createSubjectKeyId(publicKey).getKeyIdentifier();

        AlgorithmIdentifier keyEncryptionAlgorithm = null;
        byte[] encryptedKey = null;
        while (recipientInfos.hasMoreElements()) {
            ASN1Sequence seq = (ASN1Sequence) recipientInfos.nextElement();
            RecipientInfo reci = RecipientInfo.getInstance(seq);
            KeyTransRecipientInfo ktri = KeyTransRecipientInfo.getInstance(reci.getInfo());
            SubjectKeyIdentifier ski = SubjectKeyIdentifier.getInstance(ktri.getRecipientIdentifier().getId());
            if (Arrays.equals(ski.getKeyIdentifier(), keyId)) {
                keyEncryptionAlgorithm = ktri.getKeyEncryptionAlgorithm();
                encryptedKey = ktri.getEncryptedKey().getOctets();
                break;
            }
        }

        if (keyEncryptionAlgorithm == null || encryptedKey == null) {
            throw new CryptoException("The specified user is not one of the recipients of the digital envelope.");
        }

        return Pair.of(keyEncryptionAlgorithm, encryptedKey);
    }

    public static class KeyEnvelopedData {

        public final AlgorithmIdentifier keyEncryptionAlgorithm;

        public final AlgorithmIdentifier contentEncryptionAlgorithm;

        public final byte[] encryptedKey;

        public final byte[] encryptedPrivateKey;

        public KeyEnvelopedData(AlgorithmIdentifier keyEncryptionAlgorithm, AlgorithmIdentifier contentEncryptionAlgorithm, byte[] encryptedKey, byte[] encryptedPrivateKey) {
            this.keyEncryptionAlgorithm = keyEncryptionAlgorithm;
            this.contentEncryptionAlgorithm = contentEncryptionAlgorithm;
            this.encryptedKey = encryptedKey;
            this.encryptedPrivateKey = encryptedPrivateKey;
        }
    }

    /**
     * Extract Key.
     * 
     * @param publicKey the publicKey
     * @param data the data
     * @return the {@link KeyEnvelopedData}
     * @throws IOException the {@link IOException}
     */
    public static KeyEnvelopedData extractKey(final PublicKey publicKey, final byte[] data) throws IOException {
        final EnvelopedData envelopedData = EnvelopedData.getInstance(data);
        final byte[] keyId = KeyTools.createSubjectKeyId(publicKey).getKeyIdentifier();

        AlgorithmIdentifier keyEncryptionAlgorithm = null;
        byte[] encryptedKey = null;

        ASN1Set recipientInfos = envelopedData.getRecipientInfos();
        Enumeration<?> e = recipientInfos.getObjects();
        while (e.hasMoreElements()) {
            RecipientInfo recipientInfo = RecipientInfo.getInstance(e.nextElement());
            KeyTransRecipientInfo ktri = KeyTransRecipientInfo.getInstance(recipientInfo.getInfo());
            SubjectKeyIdentifier ski = SubjectKeyIdentifier.getInstance(ktri.getRecipientIdentifier().getId());
            if (Arrays.equals(ski.getKeyIdentifier(), keyId)) {
                keyEncryptionAlgorithm = ktri.getKeyEncryptionAlgorithm();
                encryptedKey = ktri.getEncryptedKey().getOctets();
                break;
            }
        }

        if (keyEncryptionAlgorithm == null || encryptedKey == null) {
            throw new CryptoException("The specified user is not one of the recipients of the digital envelope.");
        }

        AlgorithmIdentifier contentEncryptionAlgorithm = envelopedData.getEncryptedContentInfo().getContentEncryptionAlgorithm();
        byte[] encryptedPrivateKey = envelopedData.getEncryptedContentInfo().getEncryptedContent().getOctets();
        if (contentEncryptionAlgorithm == null || encryptedPrivateKey == null) {
            throw new CryptoException("The content encryption algorithm and encrypted private key cannot be null.");
        }
        
        return new KeyEnvelopedData(keyEncryptionAlgorithm, contentEncryptionAlgorithm, encryptedKey, encryptedPrivateKey);
    }
}
